// ChatSpark.js

import { BaseChatModel } from "@langchain/core/language_models/chat_models";
import crypto from "crypto";
import WebSocket from "ws";

export class ChatSpark extends BaseChatModel {
  constructor(fields) {
    super(fields);
    this.appId = fields.appId;
    this.apiKey = fields.apiKey;
    this.apiSecret = fields.apiSecret;
    this.temperature = fields.temperature ?? 0.5;
    this.maxTokens = fields.maxTokens ?? 1024;
  }

  _llmType() {
    return "spark";
  }

  async _call(messages) {
    const prompt = this.formatMessages(messages);
    return await this.callSpark(prompt);
  }

  formatMessages(messages) {
    return messages.map((msg) => {
      if (msg._getType() === "human") {
        return { role: "user", content: msg.content };
      } else {
        return { role: "assistant", content: msg.content };
      }
    });
  }

  async callSpark(messages) {
    return new Promise((resolve, reject) => {
      const host = "spark-api.xf-yun.com";
      const path = "/v4.0/chat";
      const date = new Date().toUTCString();

      const signatureOrigin = `host: ${host}\ndate: ${date}\nGET ${path} HTTP/1.1`;
      const signatureSha = crypto
        .createHmac("sha256", this.apiSecret)
        .update(signatureOrigin)
        .digest("base64");

      const authorization = `api_key="${this.apiKey}", algorithm="hmac-sha256", headers="host date request-line", signature="${signatureSha}"`;
      const authorizationBase64 = Buffer.from(authorization).toString("base64");

      const url = `wss://${host}${path}?authorization=${authorizationBase64}&date=${encodeURIComponent(
        date
      )}&host=${host}`;

      const ws = new WebSocket(url);
      let result = "";

      ws.on("open", () => {
        const payload = {
          header: {
            app_id: this.appId,
          },
          parameter: {
            chat: {
              domain: "4.0Ultra",
              temperature: this.temperature,
              max_tokens: this.maxTokens,
            },
          },
          payload: {
            message: {
              text: messages,
            },
          },
        };
        ws.send(JSON.stringify(payload));
      });

      ws.on("message", (data) => {
        const resData = JSON.parse(data.toString());
        const choices = resData.payload?.choices?.text || [];
        for (const item of choices) {
          result += item.content;
        }

        if (resData.header?.status === 2) {
          ws.close();
          resolve(result);
        }
      });

      ws.on("error", reject);
      ws.on("close", () => {});
    });
  }
}
